package dht

import (
	"log"
	"sync"

	"gopkg.in/mgo.v2/bson"
)

type FindNodeReq struct {
	Target ID `bson:"target"`
}

type FindNodeResp struct {
	Nodes []*Node `bson:"nodes"`
}

func (dht *DHT) findNodeHandler(reqb []byte) ([]byte, error) {
	req := &FindNodeReq{}
	err := bson.Unmarshal(reqb, &req)
	if err != nil {
		return nil, err
	}
	ids := dht.router.Near(req.Target, 8)
	nodes := make([]*Node, 0, len(ids))
	for _, id := range ids {
		node, ok := dht.sessions[id]
		if ok == false {
			continue
		}
		nodes = append(nodes, node)
	}
	resp := &FindNodeResp{
		Nodes: nodes,
	}
	respb, err := bson.Marshal(resp)
	if err != nil {
		return nil, err
	}
	return respb, nil
}

func (dht *DHT) findNode(addr string, target ID) ([]*Node, error) {
	req := &FindNodeReq{
		Target: target,
	}
	reqb, err := bson.Marshal(req)
	if err != nil {
		return nil, err
	}
	respb, err := dht.query(addr, FIND_NODE, reqb)
	if err != nil {
		return nil, err
	}
	resp := &FindNodeResp{}
	err = bson.Unmarshal(respb, &resp)
	if err != nil {
		return nil, err
	}
	return resp.Nodes, nil
}

func (dht *DHT) FindNode(addr string, target ID) ([]*Node, error) {
	ids := dht.router.Near(target, 8)
	nodes := map[ID]*Node{}
	visit := map[ID]struct{}{}
	for _, id := range ids {
		nodes[id] = dht.sessions[id]
	}
	// 迭代查询
	for {
		unvisit := map[ID]struct{}{}
		for _, id := range ids {
			_, ok := visit[id]
			if ok == true {
				continue
			}
			unvisit[id] = struct{}{}
		}
		if len(unvisit) == 0 {
			break
		}
		var wg sync.WaitGroup
		for id, _ := range unvisit {
			node := nodes[id]
			wg.Add(1)
			go func(node *Node) {
				defer wg.Done()
				list, err := dht.findNode(node.Addr, target)
				if err != nil {
					log.Println("ERROR:", err)
					return
				}
				dht.mu.Lock()
				defer dht.mu.Unlock()
				visit[node.ID] = struct{}{}
				for _, v := range list {
					ids = ids.Append(v.ID)
					nodes[v.ID] = v
				}
			}(node)
		}
		wg.Wait()
		ids = ids.Sort(target)
		if len(ids) > 8 {
			for _, v := range ids[8:] {
				delete(visit, v)
				delete(nodes, v)
			}
			ids = ids[:8]
		}
	}
	ret := []*Node{}
	for _, id := range ids {
		ret = append(ret, nodes[id])
	}
	return ret, nil
}
